#include <iostream>
#include <fstream>
#include <cmath>
#include "/home/mjarvis/misc/nr/ran0.cpp"
#include "/home/mjarvis/misc/nr/zbrent.cpp"
#include "Int.h"
#include "dbg.h"
using std::endl;

std::ostream* dbgout=new std::ofstream("makewwa2.debug");
//std::ostream* dbgout = &std::cout;
bool XDEBUG = false;

const char* cmbfile = "openchain.dat";

double glob_omega_m;
double glob_omega_de;
double glob_w0;
double glob_wa;

struct badh {};

double cmb_dist_integrand(double a)
{
  double lna = log(a);
  double h = glob_omega_m * exp(-3.*lna) + 
    (1.-glob_omega_m-glob_omega_de) * exp(-2.*lna) +
    glob_omega_de * exp( -3.*(1.+glob_w0+glob_wa)*lna - 3.*glob_wa*(1.-a) );
  if (h <= 0.) {
    /*
    dbg<<"Bad h found:\n";
    dbg<<"om_m = "<<glob_omega_m<<endl;
    dbg<<"om_de = "<<glob_omega_de<<endl;
    dbg<<"om_k = "<<1.-glob_omega_m-glob_omega_de<<endl;
    dbg<<"w,wa = "<<glob_w0<<"   "<<glob_wa<<endl;
    dbg<<"a = "<<a<<endl;
    dbg<<"=> hsq = "<<h<<endl;
    */
    return 1.e10;
  }
  h = sqrt(h);
  double asq = a*a;
  return 1./(h*asq);
}

double cmb_dist(double omega_m, double omega_de, double w0, double wa)
// d_a = a_cmb int( 1/H(a) a^(-2) da , a=a_cmb..1)
// H(a) = H0 sqrt( omega_m a^(-3) + omega_k a^(-2) +
//                 omega_DE exp( 3 int(1+w(a') dln(a'),a'=a..1) ) )
// With w(a) = w0 + wa(1-a), this becomes:
// H(a) = H0 sqrt( omega_m a^(-3) + omega_k a^(-2) +
//                 omega_DE exp( -3(1+w0+wa)ln(a) - 3wa(1-a) ) )
// Since we are looking for changes in this quantity for different
// parameters, I actually calculate d_a*H0/a_cmb
{
  const double acmb = 1./1100.;

  glob_omega_m = omega_m;
  glob_omega_de = omega_de;
  glob_w0 = w0;
  glob_wa = wa;
  integ::IntRegion<double> reg(acmb,1.);
  double dist = integ::int1d(std::ptr_fun(cmb_dist_integrand),reg);
  //dbg<<"dist("<<omega_m<<","<<w0<<","<<wa<<") = "<<dist<<endl;
  return dist;
}
 
const double omde_min = 0.;
const double omde_max = 2.;
const double w_min = -20.;
const double w_max = 10.;
const double wa_min = -4.;
const double wa_max = 8.;

double glob_da0;

double daroot(double w)
{
  // Want to solve for da(w) = da0
  // so root equation is da(w) - da0
  double d = cmb_dist(glob_omega_m,glob_omega_de,w,glob_wa);
  return d-glob_da0;
}

bool Bracket(double da, double& w1, double& w2)
{
  for(size_t i = 0; i<5; i++) {
    double d1 = cmb_dist(glob_omega_m,glob_omega_de,w1,glob_wa);
    double d2 = cmb_dist(glob_omega_m,glob_omega_de,w2,glob_wa);
    //dbg<<"Bracket: da("<<w1<<") = "<<d1<<",  da("<<w2<<") = "<<d2<<endl;
    if (d2 < d1) {
      double temp = w1;
      w1 = w2;
      w2 = temp;
      temp = d1;
      d1 = d2;
      d2 = temp;
    }
    // Now d1 < d2

    if (d2 < da) {
      w2 += w2-w1;
    } else if (d1 > da) {
      w1 -= w2-w1;
    } else {
      //dbg<<"Good.\n";
      return true;
    }
  }
  //dbg<<"Maxiter exceeded."<<endl;
  return false;
}

double FindW(double omega_m,double omega_de, double wa, double da)
{
  glob_da0 = da;
  glob_omega_m = omega_m;
  glob_omega_de = omega_de;
  glob_wa = wa;
  //dbg<<"da = "<<da<<endl;
  double w1 = -2., w2 = 0.;
  if (!Bracket(da,w1,w2)) {
    return -999.;
  } else {
    //dbg<<"da(minw) = "<<cmb_dist(omega_m,omega_de,w1,wa)<<endl;
    //dbg<<"da(maxw) = "<<cmb_dist(omega_m,omega_de,w2,wa)<<endl;
    double neww = zbrent(daroot,w1,w2,1.e-4);
    //dbg<<"da("<<omega_m<<","<<omega_de<<","<<neww<<","<<wa<<") = ";
    //dbg<<cmb_dist(omega_m,omega_de,neww,wa)<<endl;
    return neww;
  }
}

double range_rand(double min, double max)
{
  return min + (max-min)*ran0();
}

int main()
{
    std::ifstream fin(cmbfile);
    if (!fin) myerror("opening cmb file: ",cmbfile);

    double ombh2,omch2,omdeh2,theta,n,Z,w0,As,h,ll;
    while (fin >> ombh2 >> omch2 >> theta >> n >> Z >> omdeh2 >> As >> h >> ll) {
        double omega_m = (omch2+ombh2)/(h*h);
        double omega_de = (omdeh2)/(h*h);
        double omega_k = 1.-omega_m-omega_de;
        double chi1 = cmb_dist(omega_m,omega_de,-1.0,0.0);
        dbg<<"chi1 = "<<chi1<<endl;

        // For non flat cosmologies, the angular diameter distance is
        // Om_k > 0: sinh(sqrt(Om_k) chi) / sqrt(Om_k)
        // Om_k < 0: sin(sqrt(-Om_k) chi) / sqrt(-Om_k)
        // where chi is the comoving coordinate distance
        // The function cmb_dist calculates chi, not the real d_a.
        // chi for the above formula needs to be in units of c/H0, which
        // is how the function cmb_dist returns the value.
        double sqrtK = sqrt(fabs(omega_k));
        //dbg<<"omega_k = "<<omega_k<<endl;
        //dbg<<"sqrtK = "<<sqrtK<<endl;
        double da = omega_k = 0. ? chi1 : omega_k > 0. ? 
            sinh(sqrtK*chi1)/sqrtK : sin(sqrtK*chi1)/sqrtK;

        dbg<<"da("<<omega_m<<","<<omega_de<<",-1.0,0.0) = "<<da<<endl;
        double wa = range_rand(wa_min,wa_max);
        double newomde;
        double neww = -999.;
        double newomk, newsqrtK, newda;
        while (neww == -999.) {
            newomde = range_rand(omde_min,omde_max);
            newomk = 1.-omega_m-newomde;
            newsqrtK = sqrt(fabs(newomk));
            //dbg<<"newomk = "<<newomk<<endl;

            // Check that asin will be ok
            double dasqrtK = da*newsqrtK;
            if (newomk < 0. && fabs(dasqrtK) >= 1.) continue;

            double newchi = newomk == 0. ? da : newomk > 0. ?
                asinh(da*newsqrtK)/newsqrtK : asin(da*newsqrtK)/newsqrtK;
            //dbg<<"newchi = "<<newchi<<endl;
            neww = FindW(omega_m,newomde,wa,newchi);
        }
        // Check:
        double chi2 = cmb_dist(omega_m,newomde,neww,wa);
        //dbg<<"chi2 = "<<chi2<<endl;
        newda = newomk == 0. ? chi2 : newomk > 0. ? 
            sinh(newsqrtK*chi2)/newsqrtK : sin(newsqrtK*chi2)/newsqrtK;
        dbg<<"da("<<omega_m<<","<<newomde<<","<<neww<<","<<wa<<") = "<<newda<<endl;
#if 0
        double newwplussig = FindW(omega_m,wa,da-0.1);
        double newwminussig = FindW(omega_m,wa,da+0.1);
        double sig_w = std::abs(newwplussig - newwminussig)/2.;
        dbg<<"wa = "<<wa<<"  sig_w = "<<sig_w<<endl;
        // Use sigma_w for the likelihood to account for the phase space
        // volume at that value of wa.
        double like = sig_w;
#endif

        std::cout << ombh2 <<"  "<< omch2 <<"  "<<newomde<<"  "<<n<<"  ";
        std::cout << wa<<"  "<<neww<<"  "<<As<<"  "<<h<<"  1.0"<<std::endl;
    }
}
